package fun.ticsmyc.question;

import fun.ticsmyc.tools.ArrayTools;
import org.junit.Test;

import java.util.Arrays;

/**
 * topK问题
 *
 * 1. 排序  【略
 *  1.1 局部排序  冒泡、插排等  O（n*k）
 * 2. 维护一个堆  O（nlogk）
 * 3. BFPRT
 *
 * @author Ticsmyc
 * @date 2020-04-06 12:13
 */
public class TopK {

    public int[] partationTopK(int[] arr,int k){
        int left =0;
        int right =arr.length-1;
        k--; //转成下标
        while(left < right){
            int[] res = partation(arr,left,right);
            if( k >res[0] && k<res[1]){
                return Arrays.copyOfRange(arr,0, k+1);
            }else if( k <=res[0]){
                //左边
                right = res[0];
            }else {
                //右边
                left=res[1];
            }
        }
        return new int[]{};
    }

    public int[] partation(int[] arr,  int head, int tail){
        int left = head-1;
        int right = tail;

        //用tail所在位置的数作为分隔依据
        int partation = arr[tail];
        int cur =head;
        while(cur < right){
            if(arr[cur] < partation){
                ArrayTools.swap(arr,++left,cur++);
            }else if(arr[cur] > partation){
                ArrayTools.swap(arr,--right,cur);
            }else{
                cur++;
            }
        }
        ArrayTools.swap(arr,cur,tail);
        return new int[]{left,right+1};
    }


/**
 * 堆的方式  ：O(N + KlogN)
 *
 * 还可以进一步优化， 小顶堆的大小设为k ，剩下的和堆顶比较。比堆顶大就弹出堆顶再进堆。 O(K+（N-K）logK)
 *
 */

    /**
     * 大小为n的堆
     * @param arr
     * @param k
     * @param isTop 找最大k个还是最小k个
     */
    public void heapTopK(int[] arr,int k,boolean isTop){
        for(int i =0;i<arr.length ;++i){
            heapInsert(arr,i,isTop);
        }
        int size = arr.length;
        for(int i=0;i<k;++i){
            ArrayTools.swap(arr,0,--size);
            System.out.println(arr[size]);
            heapify(arr,0,size,isTop);
        }
    }
    public void heapTopK2(int[] arr ,int k,boolean isTop){
        //维护一个大小为k的堆
        int i =0;
        for(;i<k;++i){
            heapInsert(arr,i,!isTop);
        }
        //之后的元素都和堆顶比较
        for(; i<arr.length ;++i){
            if(isTop){
                if(arr[i] > arr[0]){
                    arr[0]=arr[i];
                    heapify(arr,0,k,false);
                }
            }else{
                if(arr[i] < arr[0]) {
                    arr[0] = arr[i];
                    heapify(arr, 0, k, true);
                }
            }
        }

        for(int j=0;j<k;++j){
            System.out.println(arr[j]);
        }

    }

    /**
     *
     * @param arr
     * @param index
     * @param isBig  大根堆还是小根堆
     */
    private void heapInsert(int[] arr, int index ,boolean isBig){
        if(isBig){
            //大顶堆
            while(arr[index]>arr[(index-1)/2]){
                ArrayTools.swap(arr,index,(index-1)/2);
                index=(index-1)/2;
            }
        }else{
            //小顶堆
            while(arr[index]<arr[(index-1)/2]){
                ArrayTools.swap(arr,index,(index-1)/2);
                index=(index-1)/2;
            }
        }

    }
    private void heapify(int[] arr, int index, int size,boolean isBig){
        if (isBig) {
            int left = 2*index+1;
            while(left<size){
                int biggest = (left+1<size && arr[left+1]>arr[left]) ? left+1 : left;
                if(arr[biggest] > arr[index]){
                    ArrayTools.swap(arr,biggest,index);
                    index = biggest;
                    left =2*index+1;
                }else{
                    break;
                }
            }
        }else{
            int left = 2*index+1;
            while(left<size){
                int smallest = (left+1<size && arr[left+1]<arr[left]) ? left+1 : left;
                if(arr[smallest] < arr[index]){
                    ArrayTools.swap(arr,smallest,index);
                    index = smallest;
                    left =2*index+1;
                }else{
                    break;
                }
            }
        }

    }

    /**
     * BFPRT方式
     */

    public int BFPRTTopK(int[] arr ,int k){
        return BFPRT.getMinKthByBFPRT(arr,k);
    }



    @Test
    public void testHeapTopK(){
        int[] arr = new int[]{4,8,9,12,5,99,89,77,1};
        //heapTopK(arr,3,false);
        //heapTopK2(arr,3,true);
//        System.out.println(BFPRTTopK(arr,3));
        ArrayTools.printArray(partationTopK(arr,5));
        ArrayTools.printArray(arr);
    }


}
